feat(profiling): real torch flame graphs - #18457

import os
from time import time

import torch
from torch import nn
from torch.profiler import ProfilerActivity


class FeedForward(nn.Module):
    def __init__(self, d_model: int, d_ff: int) -> None:
        super().__init__()
        self.w1 = nn.Linear(d_model, d_ff)
        self.act = nn.GELU()
        self.w2 = nn.Linear(d_ff, d_model)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.norm(x + self.w2(self.act(self.w1(x))))


class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int) -> None:
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.attn_norm = nn.LayerNorm(d_model)
        self.ff = FeedForward(d_model, d_ff)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        attn_out, _ = self.attn(x, x, x)
        x = self.attn_norm(x + attn_out)
        return self.ff(x)


class DeepModel(nn.Module):
    def __init__(self, d_model: int = 256, n_heads: int = 4, d_ff: int = 512, n_layers: int = 6) -> None:
        super().__init__()
        self.embed = nn.Linear(64, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers)])
        self.head = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.LayerNorm(d_model),
            nn.Linear(d_model, 16),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.embed(x)
        for block in self.blocks:
            x = block(x)
        return self.head(x)


def main() -> None:
    torch.set_num_threads(1)

    model = DeepModel()
    x = torch.randn(8, 32, 64)  # (batch, seq_len, input_dim)

    execution_time_sec = float(os.getenv("EXECUTION_TIME_SEC", "30"))
    end = time() + execution_time_sec
    with torch.no_grad():
        while time() < end:
            with torch.profiler.profile(activities=[ProfilerActivity.CPU]):
                for _ in range(10):
                    model(x)


main()